package scu.maqiang.numeric;

import java.util.Arrays;

/**
 * 行压缩稀疏矩阵类
 * @author 马强
 *
 */
public class CSRMatrix {
	//稀疏矩阵行指标
	private int[] rowIndex;
	//稀疏矩阵列指标
	private int[] columns;
	//稀疏矩阵元素值
	private double[] values;
	//稀疏矩阵行数
	private int m;
	//稀疏矩阵列数
	private int n;
	//稀疏矩阵非0元素
	private int nnz;
	//稀疏矩阵是否对称
	private boolean isSymmetric;
	
	/**
	 * 得到矩阵行数
	 * @return 矩阵行数
	 */
	public int getM() {
		return m;
	}

	/**
	 * 得到矩阵列数
	 * @return 矩阵列数
	 */
	public int getN() {
		return n;
	}

	/**
	 * 得到矩阵非零元素个数
	 * @return 矩阵非零元素个数
	 */
	public int getNNZ() {
		return nnz;
	}

	/**
	 * 得到矩阵是否对称
	 * @return 矩阵是否对称
	 */
	public boolean isSymmetric() {
		return isSymmetric;
	}

	/**
	 * 设置矩阵对称性，可将上三角矩阵设置为对称属性
	 * @param isSymmetric 矩阵对称性
	 */
	public void setSymmetric(boolean isSymmetric) {
		this.isSymmetric = isSymmetric;
	}
	
	/**
	 * 构造函数，给出矩阵行列数构造矩阵
	 * @param m 矩阵行数
	 * @param n 矩阵列数
	 */
	public CSRMatrix(int m, int n) {
		this.m = m;
		this.n = n;
		this.nnz = 0;
		rowIndex = null;
		columns = null;
		values = null;
		isSymmetric = false;
	}
	
	/**
	 * 构造函数，给出方阵阶数构造矩阵
	 * @param n 矩阵阶数
	 */
	public CSRMatrix(int n) {
		this(n, n);
	}
	
	/**
	 * 构造函数，由SRMatrix类构造矩阵
	 * @param mat SRMmatrix对象
	 */
	public CSRMatrix(SRMatrix mat) {
		if (mat.isSorted())
			mat.sort();
		isSymmetric = mat.isSymmetric();
		m = mat.getM();
		n = mat.getN();
		rowIndex = new int[m + 1];
		int colIdx = 0;
		int rowIdx;
		if (isSymmetric) {
			nnz = (mat.getNNZ() + mat.getM()) / 2;
			columns = new int[nnz];
			values = new double[nnz];
			for (int i = 0; i < m; i++) {
				rowIdx = 0;
				for (SREntry ele : mat.mat[i]) {
					if (ele.col >= i) {
						columns[colIdx] = ele.col;
						values[colIdx] = ele.val;
						colIdx++;
						rowIdx++;
					}
				}
				rowIndex[i + 1] = rowIndex[i] + rowIdx;
			}
		} else {
			nnz = mat.getNNZ();
			columns = new int[nnz];
			values = new double[nnz];
			
			for (int i = 0; i < m; i++) {
				for (SREntry ele : mat.mat[i]) {
					columns[colIdx] = ele.col;
					values[colIdx] = ele.val;
					colIdx++;
				}
				rowIndex[i + 1] = rowIndex[i] + mat.mat[i].size();
			}
		}
	}
	
	
	/**
	 * 得到矩阵行列指标下的元素值
	 * @param i 行指标
	 * @param j 列指标
	 * @return 元素值
	 */
	public double getElement(int i, int j) {
		ParamCheck.checkBounds(i, 0, m);
		ParamCheck.checkBounds(j, 0, n);
		for(int k = rowIndex[i], kEnd = rowIndex[i + 1]; k < kEnd; k++) {
			if(columns[k] == j) {
				return values[k];
			}
		}
		return 0.0;
	}
	
	/**
	 * 将SRMatrix对象转换为CSRMatrix对象, 对任意一个SRMatrix矩阵A，该静态方法将其转换为两个CSRMatrix: mat1, mat2, 
	 * 以及用一个数组diag得到该矩阵的对角元素, 如果矩阵A对称，则矩阵mat1为该对称矩阵的上三角元素，mat2为矩阵下三角元素但不包含主对角元素;
	 * 如果矩阵B不对称，则得到mat1为整个矩阵, mat2仍为原矩阵下三角元素但不包含主对角线元素构成的矩阵
	 * @param mat 一般SRMatrix矩阵
	 * @param diag 得到矩阵对角线元素
	 * @param mat1 CSRMatrix矩阵，当mat对称时，用mat上三角元素形成矩阵，并设置其为对称的; 当mat不对称时，用mat所有元素形成矩阵
	 * @param mat2 CSRMatrix矩阵，用mat下三角元素形成矩阵
	 */
	public static void SRMatrixToCSRMatrix(SRMatrix mat, double[] diag, CSRMatrix mat1, CSRMatrix mat2) {
		boolean isSymmetric = mat.isSymmetric();
		int colIdxL = 0;
		int rowIdxL = 0;
		if(isSymmetric) {
			extractDLU(mat, diag, mat1, mat2);
		}else {
			if (!mat.isSorted())
				mat.sort();
			mat1.isSymmetric = isSymmetric;
			int m = mat.getM();
			int n = mat.getN();
			mat1.m = m;
			mat1.n = n;
			mat1.rowIndex = new int[m + 1];
			int colIdx = 0;
			int nnz = mat.getNNZ();
			mat1.nnz = nnz;
			mat1.columns = new int[nnz];
			mat1.values = new double[nnz];
			
			int NNZL = mat.getNNZL();
			mat2.m = m;
			mat2.n = n;
			mat2.nnz = NNZL;
			mat2.rowIndex = new int[m + 1];
			mat2.columns = new int[NNZL];
			mat2.values = new double[NNZL];
			for (int i = 0; i < m; i++) {
				rowIdxL = 0;
				for (SREntry ele : mat.mat[i]) {
					int colj = ele.col;
					mat1.columns[colIdx] = colj;
					mat1.values[colIdx] = ele.val;
					colIdx++;
					if(colj < i) { // lower entry
						mat2.columns[colIdxL] = colj;
						mat2.values[colIdxL] = ele.val;
						rowIdxL++;
						colIdxL++;
					}
					if (colj == i) {
						diag[i] = ele.val;
					}
				}
				mat1.rowIndex[i + 1] = mat1.rowIndex[i] + mat.mat[i].size();
				mat2.rowIndex[i + 1] = mat2.rowIndex[i] + rowIdxL;
			}
		}
	}
	
	/**
	 * 对于一个SRMatrix对象的对称矩阵，将其转换为两个CSRMatrix，并得到其对角元素
	 * @param mat 对称SRMatrix矩阵
	 * @param diag 得到矩阵对角线元素
	 * @param matU 得到矩阵上三角元素，并设置矩阵对称
	 * @param matL 矩阵矩阵下三角元素，但不包含主对角线元素
	 */
	public static void extractDLU(SRMatrix mat, double[] diag, CSRMatrix matU, CSRMatrix matL) {
		int m = mat.getM();
		ParamCheck.checkEqual(diag.length, m);
		ParamCheck.checkEqual(matU.m, m);
		ParamCheck.checkEqual(matL.m, m);
		if (!mat.isSorted()) {
			mat.sort();
		}
		int NNZL = mat.getNNZL();
		int NNZU = mat.getNNZU();
		int rowIdxL, rowIdxU;
		int colIdxL = 0, colIdxU = 0;
		matU.nnz = NNZU + m;
		matU.rowIndex = new int[m + 1];
		matU.columns = new int[NNZU + m];
		matU.values = new double[NNZU + m];
		
		matL.nnz = NNZL;
		matL.rowIndex = new int[m + 1];
		matL.columns = new int[NNZL];
		matL.values = new double[NNZL];
		for(int i = 0; i < m; i++) {
			rowIdxL = 0;
			rowIdxU = 0;
			for (SREntry ele: mat.mat[i]) {
				int colj = ele.col;
				if(colj < i) { // lower entry
					matL.columns[colIdxL] = colj;
					matL.values[colIdxL] = ele.val;
					rowIdxL++;
					colIdxL++;
				}else { // upper plus diag entry
					matU.columns[colIdxU] = colj;
					matU.values[colIdxU] = ele.val;
					rowIdxU++;
					colIdxU++;					
				}
				if (colj == i) { //diag entry
					diag[i] = ele.val;
				}
			}
			matL.rowIndex[i + 1] = matL.rowIndex[i] + rowIdxL;
			matU.rowIndex[i + 1] = matU.rowIndex[i] + rowIdxU;
		}
		matU.setSymmetric(true);
	}
	

	/**
	 * 矩阵与向量乘积运算
	 * @param u 向量 
	 * @return 结果向量
	 */
	public double[] mul(double[] u) {
		double[] v = new double[m];
		mmul(u, v);
		return v;
	}
	
	/**
	 * 矩阵向量乘积运算
	 * @param u 向量
	 * @param v 结果向量
	 */
	public void mul(double[] u, double[] v) {
		ParamCheck.checkEqual(v.length, m);
		Arrays.fill(v, 0.0);
		mmul(u, v);
		//System.out.println(Arrays.toString(v));
	}
	
	/**
	 * 
	 * @param u
	 * @param v
	 */
	private void mmul(double[] u, double[] v) {
		int i, j, jEnd, colj;
		if (isSymmetric) {
			for (i = 0; i < m; i++) {
				jEnd = rowIndex[i + 1];
				for(j =  rowIndex[i]; j < jEnd; j++) {
					colj = columns[j];
					v[i] += values[j] * u[colj];
					if (i != colj) {
						v[colj] += values[j] * u[i];
					}
				}
			}
		}else {
			for (i = 0; i < m; i++) {
				jEnd = rowIndex[i + 1];
				for (j = rowIndex[i]; j < jEnd; j++) {
					v[i] += values[j] * u[columns[j]];
				}
			}
		}
	}

	/**
	 * 矩阵转置与向量乘积
	 * @param u 向量
	 * @return 结果向量
	 */
	public double[] mulT(double[] u) {
		double[] v = new double[m];
		mmulT(u, v);
		return v;
	}
	
	/**
	 * 
	 * @param u
	 * @param v
	 */
	public void mulT(double[] u, double[] v) {
		ParamCheck.checkEqual(v.length, m);
		Arrays.fill(v, 0.0);
		mmulT(u, v);
	}
	
	
	private void mmulT(double[] u, double[] v) {
		int i, j, jEnd, colj;
		if (isSymmetric) {
			for (i = 0; i < m; i++) {
				jEnd = rowIndex[i + 1];
				for(j =  rowIndex[i]; j < jEnd; j++) {
					colj = columns[j];
					v[i] += values[j] * u[colj];
					if (i != colj) {
						v[colj] += values[j] * u[i];
					}
				}
			}
		}else {
			for (i = 0; i < m; i++) {
				jEnd = rowIndex[i + 1];
				for (j = rowIndex[i]; j < jEnd; j++) {
					v[columns[j]] += values[j] * u[i];
				}
			}
		}		
	}
	
	/**
	 * 
	 * @param diag
	 * @param matU
	 * @param matL
	 * @param r0
	 * @param z0
	 * @param omega
	 */
	public static void SSORTreat(double[] diag, CSRMatrix matU, CSRMatrix matL, double[] r0, double[] z0, double omega) {
		int n = matL.getN();
		int jEnd;
		double temp;
		z0[0] = r0[0] / diag[0];
		//System.out.println(MVO.toString(diag));
		for (int i = 1; i < n; i++) {
			temp = 0.0;
			jEnd = matL.rowIndex[i + 1];
			for(int j = matL.rowIndex[i]; j < jEnd; j++) {
				temp += omega * matL.values[j] * z0[matL.columns[j]]; 
			}
			//System.out.println(r0[i] + "\t" + temp + "\t" + diag[i]);
			//System.out.println(MVO.toString(r0));
			//System.out.println(temp);
			z0[i] = (r0[i] - temp) / diag[i];
			//System.out.println(Arrays.toString(z0));
			if(Double.isNaN(z0[i])) {
				System.out.println("Singular Matrix!");
				System.exit(1);
			}
		}
		
		//System.out.println("z0 = " + MVO.toString(z0));
		for(int i = 0; i < n; i++) {
			z0[i] *= diag[i];
		}
		
		z0[n - 1] = z0[n - 1] / diag[n - 1];
		for (int i = n - 2; i >= 0; i--) {
			temp = 0.0;
			jEnd = matU.rowIndex[i + 1];
			for (int j = matU.rowIndex[i]; j < jEnd; j++) {
				if (matU.columns[j] > i) {
					temp += omega * matU.values[j] * z0[matU.columns[j]]; 
				}
			}
			z0[i] = (z0[i] - temp) / diag[i];
		}
	}

	
	public static void GSPrecondition(double[] diag, CSRMatrix matU, double[] r0, double[] z0) {
		int n = matU.getN();
		int jEnd;
		double temp;
		z0[0] = r0[0] / diag[0];
		for (int i = 1; i < n; i++) {
			temp = 0.0;
			jEnd = matU.rowIndex[i + 1];
			for(int j = matU.rowIndex[i]; j < jEnd; j++) {
				if (matU.columns[j] != i) {
					temp += matU.values[j] * z0[matU.columns[j]]; 
				}
			}
			z0[i] = (r0[i] - temp) / diag[i];
		}
	}
	
	@Override
	public String toString() {
		StringBuilder sb = new StringBuilder();
		sb.append(m + "\t" + n + "\t" + nnz + "\t" + isSymmetric +"\n");
		sb.append(Arrays.toString(rowIndex) + "\n");
		sb.append(Arrays.toString(columns) + "\n");
		sb.append(Arrays.toString(values) + "\n");
		return sb.toString();
	}
	
	/**
	 * 
	 * @return
	 */
	public String toSRFormat() {
		StringBuilder sb = new StringBuilder();
		sb.append(m + "\t" + n + "\t" + nnz + "\t" + isSymmetric +"\n");
		int jEnd, idx;
		idx = 0;
		for(int i = 0; i < m; i++) {
			jEnd = rowIndex[i + 1];
			for(int j = rowIndex[i]; j < jEnd; j++) {
				sb.append(String.format("%1$6d,%2$6d,%3$16.8f\n", i, columns[idx], values[idx]));
				idx++;
			}
		}
		return sb.toString();
	}


}
